Python 修饰器

发表于   |  

稍稍多接触一些 Python 代码，就避不开修饰器的使用。比如在 flask 中，不可避免地就会遇到 @login_required, @app.route("/") 这样的代码，这个就是 Python 的修饰器。

在 Python 中，函数是对象，函数对象可以赋值给变量，因此，通过变量可以调用这个函数： 

1
2
3
4
5
6
7
8
In [6]: def now():
   ...:     print "12:00"
   ...:

In [7]: f=now

In [8]: f()
12:00

我们知道，函数对象中有一个属性 __name__，我们可以通过它拿到函数的名称： 

1
2
3
4
5
In [9]: now.__name__
Out[9]: 'now'

In [10]: f.__name__
Out[10]: 'now'

现在我们有个 now 函数了。目前我们有一个这样的需求，我想要在代码中得到执行 now 函数的日志，例如输出运行前，运行后打印时间戳等等，但是，now 函数是不能修改的(因为项目中这种类似的代码非常多，如果改动的话引入 bug 的可能性非常大)，我们需要在代码运行期间动态给函数增加功能，这时候修饰器（decorator）就派上用场了。 我们可以用装饰器来装饰原有的函数，以实现不修改原函数却能增强原函数功能。

本质上，修饰器（decorator）是一个返回函数的高阶函数，利用装饰器，我们可以这样实现上面提到的日志函数： 

1
2
3
4
5
6
7
8
9
10
11
def log(func):
    def wrapper(*args, **kwargs):
        print 'call %s():' % func.__name__
        return func(*args, **kwargs)
    return wrapper

@log
def now():
    print u"this is now"

now()

修饰器函数以修饰的函数作为参数，并且返回一个函数，上面的代码运行结果是： 

1
2
3
$ python testdecorator.py
$ call now():
$ this is now

实际等同于： 

1
2
3
4
5
6
7
8
9
10
11
def log(func):
    def wrapper(*args, **kw):
        print 'call %s():' % func.__name__
        return func(*args, **kw)
    return wrapper

def now():
    print u"this is now"

new_now = log(now)
new_now()

这两段代码输出的结果是一样的。log 函数是一个修饰器，返回的是一个函数，因此，now()函数仍然存在，现在是同名的 now 变量指向了这个函数，用 now() 就能调用这个函数，即 log() 返回的 wrapper() 函数，wrapper 的参数定义是 *args, **kwargs， 因此可以接受任何参数，如果修饰器本身也需要参数，则需要编写一个返回 decorator 的高阶函数，如： 

1
2
3
4
5
6
7
8
9
10
11
12
13
def log(text):
    def decorator(func):
        def wrapper(*args, **kwargs):
            print '%s %s():' % (text, func.__name__)
            return func(*args, **kwargs)
        return wrapper
    return decorator

@log('run!!!')
def now():
    print u"this is now"

now()

输出结果为： 

1
2
3
$ python testdecorator.py
run!!! now():
this is now

上面的代码就相当于： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
$ cat testdecorator.py
def log(text):
    def decorator(func):
        def wrapper(*args, **kw):
            print '%s %s():' % (text, func.__name__)
            return func(*args, **kw)
        return wrapper
    return decorator

def now():
    print u"this is now"


new_now = log("run!!!")(now)
new_now()

实际上就相当于，先执行 log("run!!!")，返回 decorator 函数，此时参数是 now，输出了 run!!!now(), 返回了 wrapper，即 now 函数，通过 new_now() 的调用，输出了 this is now。

但是这里还有一个问题，前面讲到，函数有 __name__ 属性，但是经过装饰器装饰器之后的函数，__name__ 已经从原来的 now 变成了 wrapper: 

1
2
3
4
In [1]: from testdecorator import now

In [2]: now.__name__
Out[2]: 'wrapper'

这是因为返回的那个 decorator 函数名字就是 decorator， 这个问题很有解决的必要，否则有些依赖函数签名的代码就会出错，我们需要把原始代码的 __name__ 复制到 wrapper() 函数中，用 Python 内置的 functiontools.wraps 就能解决这个问题。上面的两段代码的装饰器可以分别修改为： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import functools

def log(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print 'call %s():' % func.__name__
        return func(*args, **kwargs)
    return wrapper

@log
def now():
    print u"this is now"

now()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import functools

def log(text):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            print '%s %s():' % (text, func.__name__)
            return func(*args, **kwargs)
        return wrapper
    return decorator

@log('run!!!')
def now():
    print u"this is now"

now()

写到这里，Python 装饰器的套路基本上就很清晰了。

最后，留下两个示例：

  • 编写一个 decorator，能在函数调用的前后打印出 ‘begin call’ 和 ‘end call’ 的日志。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import functools

def log(text):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            print 'begin call: ' + func.__name__
            func(*args, **kwargs)
            print 'end call: ' + func.__name__
        return wrapper
    return decorator

@log('run!!!')
def now():
    print u"this is now"

now()
  • 编写一个 decorator， 使之支持： 
1
2
3
@log
def f():
    pass

和

1
2
3
@log('execute')
def f():
    pass

这里的难点在于，当我们用 log 修饰 f() 这个函数式，实际上上面的两种情况就相当于 new_f = log(f), new_f()和 new_f = log('execute')(f), new_f(), 因此我们需要判断 log 中的参数是否是一个函数

可以利用 callable 判断参数是否是函数，答案： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import functools
def log(text):
    if callable(text):
        @functools.wraps(text)
        def wrapper(*args, **kwargs):
            print 'begin call: ' + text.__name__
            text(*args, **kwargs)
            print 'end call: ' + text.__name__
        return wrapper
    else:
        def decorator(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                print 'begin call: ' + text
                func(*args, **kwargs)
                print 'end call: ' + text
            return wrapper
        return decorator

@log
def  now1():
    print 'doing1...'

@log('text')
def now2():
    print 'doing2...'

now1()
now2()
  • 编写一个 decorator，使之支持自定义被装饰的方法的重试次数

这种需求可能比较常见，例如我们可能遇到这样的场景：某个方法是用来发送邮件的，邮件服务器不一定稳定，我们需要实现一个重试该方法的策略，连接不上服务器的时候可以重试几次，为了追求轻量级，不能使用 celery 这样的框架，这时我们就可以通过装饰器来解决：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import system
from time import sleep

def example_exc_handler(tries_remaining, exception, delay):
    @staticmethod
    def default_exc_handler(tries_remaining, exception, delay):
        """
        default exception handler
        """
        print(("[default_exc_handler] Caught '%s', %d tries remaining, "
                      "sleeping for %s seconds") % (exception, tries_remaining, delay))

def retries(max_tries, delay=1, backoff=2, exceptions=(Exception,), hook=None):
    def dec(func):
        def f2(*args, **kwargs):
            mydelay = delay
            tries = range(max_tries)
            tries.reverse()
            for tries_remaining in tries:
                try:
                   return func(*args, **kwargs)
                except exceptions as e:
                    if tries_remaining > 0:
                        if hook is not None:
                            hook(tries_remaining, e, mydelay)
                        sleep(mydelay)
                        mydelay = mydelay * backoff
                    else:
                        raise
                else:
                    break
        return f2
    return dec

一份简单的测试代码是：

1
2
3
4
5
6
7
8
if __name__ == "__main__":

   # @retries(max_tries=2, delay=1, backoff=2)
   def test_print(*args, **kwargs):
       print("test")
       raise Exception

   test_print()

参考：https://gist.github.com/n1ywb/2570004